from django import forms
from htstool.models import *

from util import *
from stat import *
from rankprod import rp
from scipy.stats import pearsonr
from numpy import log
import pandas as pd

import json
import itertools
import sys
import zipfile
import cStringIO



# forms for the importing wizard to create experiment and plates
class ExpForm(forms.Form):
	expname=forms.CharField()
	group=forms.CharField()
	screen=forms.CharField()
	screentype=forms.ChoiceField(widget=forms.RadioSelect, choices=[('rnai','RNAi'), ('cmpd','Compound')])
	library=forms.CharField()
	description=forms.CharField()
	
	# save experiment names and desc
	def save(self):
		exp=Experiment(**self.cleaned_data)
		exp.save()
	
		return exp

class PlateSettingForm(forms.Form):
	user = forms.EmailField()
	treatment = forms.CharField(widget=forms.Textarea)
	condition = forms.CharField()
	
	# this field is a json string of ctrls selected by a user
	ctrls = forms.CharField(widget=forms.HiddenInput())

	size = forms.CharField(widget=forms.HiddenInput())

	ctrl_types = forms.ChoiceField(widget=forms.RadioSelect, choices=Control.CTRL_CHOICES)
	
	# save setting for the experiment
	def save(self, exp):
		setting = PlateSetting(**self.cleaned_data)
		setting.save()
		exp.platesetting = setting
		exp.save()


class FileUpForm(forms.Form):
	rawfile = forms.FileField()


class FileConfirmForm(forms.Form):
	confirm = forms.BooleanField()
	barcode = forms.CharField(widget=forms.HiddenInput())
	
	def get_plates(self, f_uploaded, p_size):
		if f_uploaded.content_type=="application/zip":
			zpdata = cStringIO.StringIO(f_uploaded.read())
			zp=zipfile.ZipFile(zpdata)

			# get one plate data and feed back to user to confirm
                       	fname = zp.namelist()[0]
			if p_size=='96':
				dim = (8, 12)
			else:
				dim = (16, 24)
			
			data = cStringIO.StringIO(zp.read(fname))

			# read in the raw data with Envision format 
			if fname.endswith('.xls') or fname.endswith('.xlsx'):
				dfs = read_envision_format(data, dim, excel=True)
			else:
				dfs = read_envision_format(data, dim)

			return dfs
	


		

class BarcodeForm(forms.Form):
	cond_str = forms.CharField()
	concen_str = forms.CharField()
	repl_str = forms.CharField()
	cell_str = forms.CharField()
	lib_part_str = forms.CharField()


class AnnoUpForm(forms.Form):
	annofile = forms.FileField(required=False)










class NormalisationForm(forms.Form):
	def __init__(self, *args, **kwargs):
		super(NormalisationForm, self).__init__(*args, **kwargs)
		raw_g=PlateGroup.objects(plate_type=PlateGroup.RAW)
		summ_g=PlateGroup.objects(plate_type=PlateGroup.SUMM)
		all_g=[p for p in raw_g] + [p for p in summ_g]
		self.fields['plategroupid']=forms.ChoiceField(choices=[(pg.id, pg.groupname+pg.groupdesc) if pg.groupdesc is not None else (pg.id, pg.groupname) for pg in all_g])
	
	norm_method=forms.ChoiceField(choices=NormPlate.NORM_CHOICES)

	def normalise(self, norm_method, groupid):
		"""
		normalise plates by the specified norm method
		"""
		print groupid
		plates=getplatesbygrpid(groupid)
		plgroup=PlateGroup.objects.get(id=groupid)
		exp=Experiment.objects.get(plategroups=plgroup)
		
		newplategroup=PlateGroup(groupname=plgroup.groupname, 
				plate_type=PlateGroup.NORM, 
				groupdesc=plgroup.groupdesc+norm_method,
				platecommonattr=plgroup.platecommonattr)
		
		# TODO -- make copy object attr util
		# Per plate z score normalisation: 
		# zscore = (PE - PE_mean)/PE_stdev
		if norm_method==NormPlate.PERPLATE:

			for plate in plates:
				df=plate.data.dataframe
				z_per_plate=(df - df.values.mean())/df.values.std()

				df_dict=z_per_plate.to_dict()
				dataset=Dataset(data=df_dict)
				dataset.save()

				normplate=NormPlate(normmethod=NormPlate.PERPLATE, plate_normalised=plate, data=dataset)
				normplate=cp_from_plate(normplate, plate)
				
				newplategroup.plates.append(normplate)
			newplategroup.save()

		# Per plate robust Z score normalisation	
		# rzscore = (PE - PE_median)/MAD(PE)
		elif norm_method==NormPlate.ROBUSTZ:
			
			for plate in plates:
				df=plate.data.dataframe
				mad = MAD(df.values)
				rz_per_plate=(df - np.median(df.values))/mad
				
				df_dict=rz_per_plate.to_dict()
				dataset=Dataset(data=df_dict)
				dataset.save()

				normplate=NormPlate(normmethod=NormPlate.PERPLATE,plate_normalised=plate, data=dataset)
				normplate=cp_from_plate(normplate, plate)
				
				newplategroup.plates.append(normplate)
			newplategroup.save()

		# WARNING!!!!! concentration need to be added!!!
		# B score normalisation	
		# B score = MedianPolish(data)/MAD(data)
		elif norm_method==NormPlate.BSCORE:
			
			for plate in plates:
				df=plate.data.dataframe
				mp = MedianPolish(df.values)
				ge, ce, re , resid, tbl_org =  mp.median_polish(4)
				
				bscore=resid/MAD(resid)
				bscore_df=pd.DataFrame(bscore, index=df.index, columns=df.columns)

				df_dict=bscore_df.to_dict()
				dataset=Dataset(data=df_dict)
				dataset.save()

				normplate=NormPlate(normmethod=NormPlate.PERPLATE,plate_normalised=plate, data=dataset)
				normplate=cp_from_plate(normplate, plate)
				
				newplategroup.plates.append(normplate)
			newplategroup.save()


		# PERPLATE_NT normalisation: 
		# nt_z score = (PE - PE_NT_mean)/PE_NT_stdev
		elif norm_method==NormPlate.PERPLATE_UNTREATED:

			for plate in plates:
				unt_ctrls=plate.getUntreatedcontrols()
				UNT_letter= [unt.letter for unt in unt_ctrls]
				UNT_number= [unt.num for unt in unt_ctrls]
				unt_df=plate.data.dataframe[UNT_number].loc[UNT_letter]

				df=plate.data.dataframe
				z_per_plate=(df - unt_df.values.mean())/unt_df.values.std()

				df_dict=z_per_plate.to_dict()
				dataset=Dataset(data=df_dict)
				dataset.save()

				normplate=NormPlate(normmethod=NormPlate.PERPLATE,plate_normalised=plate, data=dataset)
				normplate=cp_from_plate(normplate, plate)
				
				newplategroup.plates.append(normplate)
			newplategroup.save()
		

		# PERPLATE_NT normalisation: 
		# nt_z score = (PE - PE_NT_mean)/PE_NT_stdev
		elif norm_method==NormPlate.PERPLATE_NT:

			for plate in plates:
				nt_ctrls=plate.getNTcontrols()
				NT_letter= [nt.letter for nt in nt_ctrls]
				NT_number= [nt.num for nt in nt_ctrls]
				nt_df=plate.data.dataframe[NT_number].loc[NT_letter]

				df=plate.data.dataframe
				z_per_plate=(df - nt_df.values.mean())/nt_df.values.std()

				df_dict=z_per_plate.to_dict()
				dataset=Dataset(data=df_dict)
				dataset.save()

				normplate=NormPlate(normmethod=NormPlate.PERPLATE,plate_normalised=plate, data=dataset)
				normplate=cp_from_plate(normplate, plate)
				
				newplategroup.plates.append(normplate)
			newplategroup.save()
		
		exp=Experiment.objects.get(plategroups=plgroup)
		exp.plategroups.append(newplategroup)
		exp.save()


class SummarisationForm(forms.Form):
	def __init__(self, *args, **kwargs):
		super(SummarisationForm, self).__init__(*args, **kwargs)
		self.fields['nondruggroupid']=forms.ChoiceField(choices=[(pg.id, pg.groupname) for pg in PlateGroup.objects])
		self.fields['druggroupid']=forms.ChoiceField(choices=[(pg.id, pg.groupname) for pg in PlateGroup.objects])

	summ_method=forms.ChoiceField(choices=SummPlate.SUMM_CHOICES)
	
	def summarise(self, summ_method, non_drug_grpid, drug_grpid):
		"""
		summarise plates by the specified summ method
		"""
		non_drug_plates=getplatesbygrpid(non_drug_grpid)
		drug_plates=getplatesbygrpid(drug_grpid)
		nondruggroup=PlateGroup.objects.get(id=non_drug_grpid)
		druggroup=PlateGroup.objects.get(id=drug_grpid)
		newcond=druggroup.platecommonattr['condition']+' - '+nondruggroup.platecommonattr['condition']
		newplategroup=PlateGroup(groupname='summarised:'+summ_method+' '+newcond, plate_type=PlateGroup.SUMM, platecommonattr={'condition':newcond})
		
		newplategroup.groupdesc=druggroup.groupdesc+summ_method+nondruggroup.groupdesc

		# SUBTRACT summarisation: summ = drug - nondrug
		if summ_method==SummPlate.SUBTRACT:
			
			#TODO -- sorting based on plate name
			non_pls_sorted=sorted(non_drug_plates, key=lambda p: sort_pls(p))
			drug_pls_sorted=sorted(drug_plates, key=lambda p: sort_pls(p))

			for nondrugplate, drugplate in zip(non_pls_sorted, drug_pls_sorted):
				
				if nondrugplate.replicate!=drugplate.replicate:
					for p1,p2 in  zip([i.platename for i in non_pls_sorted], [i.platename for i in drug_pls_sorted]):
						print p1, p2
					raise Exception('drug vs no drug plates not matching')
				else:
					nondrug=nondrugplate.data.dataframe
					drug=drugplate.data.dataframe
					
					# drug - nondrug plate or irrad - nonrad
					# minus value for killing effect
					summ_plate = drug - nondrug

					df_dict=summ_plate.to_dict()
					dataset=Dataset(data=df_dict)
					dataset.save()

					summplate=SummPlate(summmethod=SummPlate.SUBTRACT, plates_summed=[nondrugplate, drugplate], data=dataset)
					summplate=cp_from_plate(summplate, drugplate)

					newplategroup.plates.append(summplate)
		
		elif summ_method==SummPlate.RATIO:

			non_pls_sorted=sorted(non_drug_plates, key=lambda p: sort_pls(p))
			drug_pls_sorted=sorted(drug_plates, key=lambda p: sort_pls(p))

			for nondrugplate, drugplate in zip(non_pls_sorted, drug_pls_sorted):
				
				if nondrugplate.replicate!=drugplate.replicate:
					for p1,p2 in  zip([i.platename for i in non_pls_sorted], [i.platename for i in drug_pls_sorted]):
						print p1, p2
					raise Exception('drug vs no drug plates not matching')
				else:
					nondrug=nondrugplate.data.dataframe
					drug=drugplate.data.dataframe
					
					# drug / nondrug plate or irrad / nonrad
					#summ_plate = log(drug / nondrug)
					summ_plate = (drug / nondrug)

					df_dict=summ_plate.to_dict()
					dataset=Dataset(data=df_dict)
					dataset.save()

					summplate=SummPlate(summmethod=SummPlate.RATIO, plates_summed=[nondrugplate, drugplate], data=dataset)
					summplate=cp_from_plate(summplate, drugplate)
					
					newplategroup.plates.append(summplate)
		
		newplategroup.save()
		exp=Experiment.objects.get(plategroups=druggroup)
		exp.plategroups.append(newplategroup)
		exp.save()



####### Form and actions to create box plots for groups of plates ########
class BoxGroupPlates(forms.Form):
	start_plate = forms.IntegerField(initial=1)
	end_plate = forms.IntegerField(initial=10)
	#groupid = forms.ChoiceField(choices=[(pg.id, pg.groupname) for pg in PlateGroup.objects])
	groupid = forms.CharField()
	
	def get_box_json(self, grpid, num_plates):
		start, end=num_plates

		plategroup = PlateGroup.objects.get(id=grpid)

		df=convert_to_one_col_df(plategroup.plates[0].data.dataframe, plategroup.plates[0].platename
				+'-'+str(plategroup.plates[0].replicate))
		labels=df.index.tolist()
		data=[]
	
		# NB: groups is misleading - it means plates grouped by the same name
		groups = groupby_platename(plategroup.plates[start:end+1])
		
		for group in groups:
			ctrls=None
			if len(group)==2:
				p1, p2 = group
				arr1=p1.data.dataframe.values.ravel()
				arr2=p2.data.dataframe.values.ravel()

				arr1, arr2=rm_nan(arr1, arr2)

				corr, pval=pearsonr(arr1, arr2)

				plate=[(p1.platename+'-'+str(p1.replicate), arr1.tolist(), str(p1.id), corr, pval,), 
					(p2.platename+'-'+str(p2.replicate), arr2.tolist(), str(p2.id), corr, pval,)]
				data.append(plate)
				ctrls=p1.getAllcontrols()
			else:
				plate=[(p.platename+'-'+str(p.replicate), p.data.dataframe.values.ravel().tolist(), str(p.id)) for p in group]
				data.append(plate)
				ctrls=p.getAllcontrols()


		json_data=json.dumps(data)
		json_labels=json.dumps(labels)
		return (json_data, json_labels, ctrls, plategroup.groupname)



####### Form and actions to create errbar plots for groups of plates ########
class ErrbarGroupPlates(forms.Form):
	groups = forms.MultipleChoiceField(choices=[(pg.id, pg.groupname+pg.groupdesc) if pg.groupdesc is not None else (pg.id, pg.groupname)for pg in PlateGroup.objects])
	
	def get_mean_json(self, grpids):
		#def selected(g):
		#	return g.id in grpids

		#filtered = filter(selected, groups)
		groups=[PlateGroup.objects.get(id=grpid) for grpid in grpids]
		all_objects=[self.get_formated_data(g) for g in groups]
		err_data=json.dumps(list(itertools.chain(*all_objects)))

		return err_data

	
	def get_formated_data(self, group):

		def mask_nan(arr):
			m_arr=np.ma.masked_array(arr, np.isnan(arr))
			return m_arr

		objs=[]
		for p in group.plates:
			nt_ctrls=p.getNTcontrols()
			NT_letter= [nt.letter for nt in nt_ctrls]
			NT_number= [nt.num for nt in nt_ctrls]
			nt_df=p.data.dataframe[NT_number].loc[NT_letter]
			
			objs.append({'Platename':p.platename+' R'+str(p.replicate), 'Condition':p.condition, 'NTmean':mask_nan(nt_df.values).mean(), 'NTstd':mask_nan(nt_df.values).std(),
			'Platemean':mask_nan(p.data.dataframe.values).mean(), 'Platestd':mask_nan(p.data.dataframe.values).std(), 'URL':"/p/"+str(p.id),})

		return objs
	

####### Forms for actions and analysis such as RankProd ########
class SelectActionForm(forms.Form):
	ACT_CHOICES = NormPlate.NORM_CHOICES + SummPlate.SUMM_CHOICES + RankProdPlate.RP_CHOICES
	action = forms.ChoiceField(widget=forms.RadioSelect, choices=ACT_CHOICES)

class RPsettingForm(forms.Form):
	perm_num = forms.IntegerField(initial=100)
	merge_plates = forms.BooleanField(required=False)
	def act(self, plate_objs):
		if self.cleaned_data['merge_plates'] is True:
			try:
				perm = self.cleaned_data['perm_num']
				pl_ids = [p_obj['plateid'] for p_obj in plate_objs]
				plates = Plate.objects(id__in=pl_ids)
			
				# only take replica 1 plates 
				#rep1_pl_ids = [str(p.id) for p in plates if p.replicate == 1]
				#a,b=plates[0].dimension
				#n=a*b
				#plateids_index = reduce(lambda x,y:x+y, [[i]*n for i in self.plateids])
				
				superplate, anno = merge_plates(plates)
				#anno_df=pd.DataFrame(data=anno, index=superplate.index, columns=['Annotation'])
				
				newplategroup=PlateGroup(groupname='RankProd - merged plates: '	, plate_type=PlateGroup.ANALYSIS)
			
				RP, pfp, pval = rp(superplate.values, perm)
				
				RP_df=pd.DataFrame(data=RP, index=superplate.index, columns=['RankProd'])
				pfp_df=pd.DataFrame(data=pfp, index=superplate.index, columns=['PFP'])
				pval_df=pd.DataFrame(data=pval, index=superplate.index, columns=['P-val'])
				anno_df=pd.DataFrame(data=anno, index=superplate.index, columns=['Annotation'])

				dataset=Dataset(data=superplate.to_dict())
				dataset.save()

				rp_dataset=Dataset(data=RP_df.to_dict())
				rp_dataset.save()

				pfp_dataset=Dataset(data=pfp_df.to_dict())
				pfp_dataset.save()

				pval_dataset=Dataset(data=pval_df.to_dict())
				pval_dataset.save()

				anno_dataset=Dataset(data=anno_df.to_dict())
				anno_dataset.save()

				rp_plate=RPSuperPlateSummary(data=dataset,
				rankprod=rp_dataset, pfp=pfp_dataset, pval=pval_dataset, annotation=anno_dataset,
				plateids=pl_ids)

				rp_plate.save()
				
				newplategroup.plates.append(rp_plate)
				newplategroup.save()
				

				grp = PlateGroup.objects.get(plates__in=pl_ids)
				exp=Experiment.objects.get(plategroups=grp)
				exp.plategroups.append(newplategroup)
				exp.save()

			except:
				print "Creating SuperPlate error:", sys.exc_info()
				raise
		else:
			# todo - per plate rank prod
			print 'per plate'
			pl_ids = [p_obj['plateid'] for p_obj in plate_objs]
			print pl_ids


class RankProd(forms.Form):
	def __init__(self, *args, **kwargs):
		super(RankProd, self).__init__(*args, **kwargs)
		self.fields['groupid']=forms.ChoiceField(choices=[(pg.id, pg.groupname) for pg in PlateGroup.objects])
		
	
	def rankprod(self, groupid, rule=None, reverse=False):
		grp=PlateGroup.objects.get(id=groupid)
		plates=grp.plates
		
		newplategroup=PlateGroup(groupname='RankProd: ', plate_type=PlateGroup.ANALYSIS, 
				groupdesc=grp.groupdesc,
				platecommonattr={'condition':rule})

		# Group by platename - repls is a list of [rep1, rep2, rep3...] of the same plate
		repls=groupby_platename(plates)
		
		# TODO - add IF to separate superplate and perplate rankprod
		for plate_reps in repls:
			dfs=[]
			for plate_rep in plate_reps:
				df=convert_to_one_col_df(plate_rep.data.dataframe, 'Replicate'+str(plate_rep.replicate))
				dfs.append(df)
			
			# do the rank product analysis, 
			# the returned dataframe contains RankProd, pfp, pval and data
			df_rp_with_data = do_rankprod(dfs, reverse=reverse)
			
			rankprod=convert_to_plate_df(df_rp_with_data['RankProd'])
			pfp=convert_to_plate_df(df_rp_with_data['pfp'])
			pval=convert_to_plate_df(df_rp_with_data['pval'])


			rp_dataset=Dataset(data=rankprod.to_dict())
			rp_dataset.save()

			pfp_dataset=Dataset(data=pfp.to_dict())
			pfp_dataset.save()

			pval_dataset=Dataset(data=pval.to_dict())
			pval_dataset.save()
			
			if rule is None:
				rule='single cell line, single concentration'
			rp_plate=RankProdPlate(ranking_rule=rule,
			rankprod=rp_dataset, pfp=pfp_dataset, pval=pval_dataset, replicate=plate_reps)
			
			attrnames=['cellline', 'platename', 'controls',  
			'concentration', 'dimension', 'condition']
			cp_plate_attr(plate_rep, rp_plate, attrnames)
			
			rp_plate.save()
			newplategroup.plates.append(rp_plate)

			df_rp_with_data.to_csv(plate_rep.platename+rule+'.csv')


		newplategroup.save()
		exp=Experiment.objects.get(plategroups=grp)
		exp.plategroups.append(newplategroup)
		exp.save()




ACTION_FORM_MAP = {'rankprod': RPsettingForm}
			
